import datetime
import json

import akshare as ak
import pandas as pd
import numpy as np
import matplotlib
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.sql import text
from decouple import config
import time


class StockStrategy:
    # 指数最大跌幅
    MAX_FALL_BASIC_PRICE_LIMIT = -1.0

    # 涨幅区间
    MAX_PRICE_LIMIT = 5
    MIN_PRICE_LIMIT = 3

    # 振幅区间
    MAX_AMPLITUDE_PRICE_LIMIT = 5
    MIN_AMPLITUDE_PRICE_LIMIT = 3

    # 流通市值区间
    MAX_CIRCULATING_MARKET_VALUE = 50 * 1e9
    MIN_CIRCULATING_MARKET_VALUE = 10 * 1e9

    # 换手率
    TURNOVER_RATE = 3

    # 涨停
    LIMIT_UP_PAST_DAY = 20

    # 量比
    VOLUME_RATIO = 1.2

    # 大盘指数exponent，沪深300 = 000300，上证50=000016
    BASIC_EXPONENT = "000300"

    # 选股的交易日
    BASIC_TRADE_DATE = "20240520"

    def __init__(self):
        self.__init_db__()

    def __init_db__(self):
        self.host = config("MYSQL_HOST")
        self.port = config("MYSQL_PORT")
        self.user = config("MYSQL_USER")
        self.password = config("MYSQL_PASSWORD")
        try:
            self.engine = create_engine('mysql+pymysql://{username}:{password}@{host}:{port}/stock?charset=utf8'
                                        .format(username=self.user, password=self.password, host=self.host, port=self.port))
        except:
            print("something wrong!")

    def __delete__(self, table_name: str, where: str):
        # 创建会话
        Session = sessionmaker(bind=self.engine)
        session = Session()

        # 执行SQL语句
        sql = "delete from {table}".format(table=table_name)
        if where is not None:
            sql = sql + " where " + where
        # 提交会话
        result = session.execute(text(sql))
        session.commit()

        # 关闭会话
        session.close()

    def save_to_mysql(self, data: pd.DataFrame):
        data.to_sql("trend_analysis", con=self.engine, if_exists='append', index=False)

    def is_current_strong(self, stock_id: str):
        # get trade date, past 14 days
        # trade_date = datetime.datetime.strftime(self.get_trade_date(), "%Y-%m-%d")
        day_0 = datetime.datetime.now()
        day_14 = day_0 - datetime.timedelta(days=14)
        # 过往14天交易日
        date_list = self.get_date_range(day_14, day_0)
        # 所有交易日
        trade_date_list = self.get_trade_date_list()

        actual_trade_date_list = set(date_list).intersection(set(trade_date_list))

        strong_stocks = []
        for trade_date in actual_trade_date_list:
            strong_stocks_df = ak.stock_zt_pool_strong_em(trade_date.strftime("%Y%m%d"))
            for index, row in strong_stocks_df.iterrows():
                strong_stocks.append(row["代码"])

        return strong_stocks.__contains__(stock_id)

    def get_date_range(self, start_date, end_date):
        date_list = [start_date + datetime.timedelta(days=x) for x in range((end_date - start_date).days + 1)]
        return [x.date() for x in date_list]

    def get_trade_date_list(self):
        trade_date_list = []
        trade_date_df = ak.tool_trade_date_hist_sina()
        for index, row in trade_date_df.iterrows():
            trade_date_list.append(row["trade_date"])

        return trade_date_list

    def run(self):
        basic_exponent_df = ak.stock_zh_index_spot_em()

        for index, row in basic_exponent_df.iterrows():
            # print(index, row)
            if row["代码"] == self.BASIC_EXPONENT:
                if row["涨跌幅"] <= self.MAX_FALL_BASIC_PRICE_LIMIT:
                    print("code: {code}, " + "Price limit：{price_limit}".format(code=row["代码"], price_limit=row["涨跌幅"]))
                    return
            else:
                continue

        # 执行策略
        # 1) 振幅在5%以内
        # 2) 流通市值在200亿以下
        # 3) 换手率3%以上
        # 4) 20天内有涨停板趋势历史
        # 5) 量比在1.2以上
        stocks = ak.stock_zh_a_spot_em()
        match_stocks = []
        for index, stock in stocks.iterrows():
            # stock_act = ak.stock_zh_a_hist(symbol=stock["代码"], period="daily", start_date=self.BASIC_TRADE_DATE,
            #                                end_date=self.BASIC_TRADE_DATE, adjust="qfq")
            # 1：日期 2：股票代码 3：开盘 4：收盘 5：最高 6：最低 7：成交量 8：成交额 9：振幅 10：涨跌幅 11：涨跌额 12：换手率
            if self.MIN_PRICE_LIMIT <= stock["涨跌幅"] \
                    and self.MIN_CIRCULATING_MARKET_VALUE <= stock["流通市值"] <= self.MAX_CIRCULATING_MARKET_VALUE \
                    and stock["换手率"] >= self.TURNOVER_RATE \
                    and stock["量比"] >= self.VOLUME_RATIO \
                    and stock["振幅"] <= self.MAX_AMPLITUDE_PRICE_LIMIT:
                match_stocks.append({stock["代码"], stock["名称"]})

                # 构建数据集
                day_0 = datetime.datetime.now().date()
                # 所有交易日
                trade_date_list = self.get_trade_date_list()
                trade_date = min(trade_date_list, key=lambda date: abs(day_0 - date))
                data = [stock["代码"], stock["名称"], stock["最新价"], trade_date.strftime("%Y%m%d"), json.dumps(stock.to_dict())]
                df = pd.DataFrame(columns=["stock_id", "stock_name",
                                           "closed_price", "trade_date", "info"], data=[data])
                print("data: {data}".format(data=str(data)))
                # delete the old record
                self.__delete__(table_name="trend_analysis", where="stock_id='{stock_id}' and trade_date='{trade_date}'"
                                .format(stock_id=stock["代码"], trade_date=trade_date.strftime("%Y%m%d")))
                # save the latest record
                self.save_to_mysql(df)

        print("Match Stocks: {result}".format(result=str(match_stocks)))


if __name__ == '__main__':
    stock_collector = StockStrategy()
    stock_collector.run()
